# main.py
import os, glob, sys
import fitz  # PyMuPDF
from fix_torch import fix_torch_dll_path
fix_torch_dll_path()
import traceback
from utils.pdf_utils import extract_pages, extract_text_words
from utils.ner_utils import detect_sensitive
from utils.desense_utils import apply_text_desense, apply_image_desense
from utils.logger import logger
from config import DESENSE_METHODS
from utils.rapid_ocr_extract import RapidOCRExtractor
import io
sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8', errors='replace')
sys.stderr = io.TextIOWrapper(sys.stderr.buffer, encoding='utf-8', errors='replace')
def classify_page_type(page):
    """
    判断单个页面是图片型还是数字型
    基于文本提取能力的简单检测方法
    
    Args:
        page: PDF页面对象
        
    Returns:
        str: 'image' 表示图片型页面, 'text' 表示数字型页面
    """
    try:
        # 尝试直接提取文本
        chars = extract_text_words(page)
        
        # 判断标准：如果能提取到足够多的文本字符（>10个），认为是数字型页面
        if chars and len(chars) > 10:
            return 'text'
        else:
            return 'image'
            
    except Exception as e:
        logger.error(f"页面类型检测失败: {e}")
        return 'image'  # 默认按图片型处理



def process_pdf(path, output_dir, method='mask', mask_color=None, mosaic_size=None, replace_char=None):
    """
    统一处理PDF文件的主函数 - 按页分别处理
    """
    logger.info(f"开始处理PDF文件: {os.path.basename(path)}")
    
    doc = extract_pages(path)
    ocr_extractor = RapidOCRExtractor()
    ocr_results = None  # 只调用一次OCR
    
    for idx in range(len(doc)):
        page = doc[idx]
        
        # 对每一页进行类型检测
        page_type = classify_page_type(page)
        logger.info(f"第 {idx+1} 页类型检测结果: {page_type}型")
        
        if page_type == 'text':
            # 数字型页面处理
            logger.info(f"{os.path.basename(path)} page#{idx+1}: 数字型页面，直接提取文本")
            chars = extract_text_words(page)
            
            if chars and len(chars) >= 10:
                hits = detect_sensitive(chars)
                # logger.info(f"chars: {chars}")
                logger.info(f"{os.path.basename(path)} page#{idx+1}: found {len(hits)} items")
                
                if method in ['replace']:
                    logger.info("replace")
                    apply_text_desense(page, hits, method=method, replace_char=replace_char)
                else:
                    logger.info("blur")
                    apply_image_desense(page, hits, method=method, mask_color=mask_color, mosaic_size=mosaic_size, replace_char=replace_char)
            else:
                logger.warning(f"{os.path.basename(path)} page#{idx+1}: 文本提取失败，跳过处理")
        else:
            # 图片型页面处理
            logger.info(f"{os.path.basename(path)} page#{idx+1}: 图片型页面，使用RapidOCR处理")
            
            if ocr_results is None:
                ocr_results = ocr_extractor.extract_from_pdf(page)
            
            page_result = ocr_results
            if page_result and page_result['chars']:
                chars = [
                    (char['bbox'][0], char['bbox'][1], char['bbox'][2], char['bbox'][3], char['char'])
                    for char in page_result['chars']
                ]
                logger.info(f"RapidOCR识别成功，共识别 {len(chars)} 个字符")
                
                hits = detect_sensitive(chars)
                logger.info(f"{os.path.basename(path)} page#{idx+1}: found {len(hits)} items")
                
                if method in ['replace']:
                    logger.info("replace")
                    apply_text_desense(page, hits, method=method, replace_char=replace_char)
                else:
                    logger.info("blur")
                    apply_image_desense(page, hits, method=method, mask_color=mask_color, mosaic_size=mosaic_size, replace_char=replace_char)
            else:
                logger.warning(f"RapidOCR识别失败，无法处理该页面")
    
    # 保存处理后的PDF
    out = os.path.join(output_dir, os.path.basename(path))
    doc.save(out)
    logger.info(f"PDF处理完成，保存到: {out}")

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="PDF脱敏处理工具 - 自动识别PDF类型并分别处理")
    parser.add_argument("--input_glob", default="./data/公证书.pdf", help="待脱敏文件或目录通配符，例如 './data/*.pdf'")
    parser.add_argument("--out", default="./out", help="输出目录")
    parser.add_argument("--method", default="mask", choices=DESENSE_METHODS, help="脱敏方法")
    parser.add_argument("--mask_color", default=None, help="遮罩颜色 (如 #000000 为黑色)")
    parser.add_argument("--mosaic_size", type=int, default=None, help="马赛克块大小 (像素)")
    parser.add_argument("--replace_char", default=None, help="替换字符 (如 X, *, #)")
    args = parser.parse_args()

    # 确保输出目录存在
    os.makedirs(args.out, exist_ok=True)
    
    # 获取所有PDF文件
    pdf_files = glob.glob(args.input_glob)
    if not pdf_files:
        logger.warning(f"在 {args.input_glob} 中未找到PDF文件")
        sys.exit(1)
    
    logger.info(f"找到 {len(pdf_files)} 个PDF文件待处理")
    
    # 处理每个PDF文件
    success_count = 0
    for pdf_file in pdf_files:
        try:
            logger.info(f"\n{'='*50}")
            logger.info(f"处理文件: {pdf_file}")
            logger.info(f"{'='*50}")
            
            process_pdf(
                pdf_file, 
                args.out, 
                method=args.method, 
                mask_color=args.mask_color, 
                mosaic_size=args.mosaic_size,
                replace_char=args.replace_char
            )
            success_count += 1
            
        except Exception as e:
            # 1. 异常类型
            exc_type = type(e).__name__
            # 2. 异常值（repr 可以看到更多细节）
            exc_value = repr(e)
            # 3. 完整堆栈跟踪
            full_tb = traceback.format_exc()

            logger.error(
                "====== PDF 处理失败 ======\n"
                "文件路径   : %s\n"
                "异常类型   : %s\n"
                "异常信息   : %s\n"
                "堆栈跟踪   :\n%s"
                "=========================",
                pdf_file,
                exc_type,
                exc_value,
                full_tb
            )
    
    logger.info(f"\n{'='*50}")
    logger.info(f"处理完成！成功处理 {success_count}/{len(pdf_files)} 个文件")
    logger.info(f"输出目录: {args.out}")
    logger.info(f"{'='*50}")
